#include <src/pcps.h>
#include <src/input.h>
#include <src/hamiltonian.h>
#include <src/wavefunction.h>
#include <src/vmc_walker.h>
#include <src/energy_accumulator.h>
#include <src/restricted_pairing_function.h>
#include <src/matrix_math.h>
#include <src/timing.h>
#include <src/archive.h>

template <class S> void exact_sampling(const pcps::Input & userinp,
                                       pcps::Hamiltonian<S> & ham,
                                       pcps::Wavefunction<S> & wfn) {

  // create the walker
  pcps::VMC_Walker<S> walker(userinp);

  // set the guiding function
  walker.gf() = wfn.overall_coefficient();

  // prepare accumulator
  boost::shared_ptr<pcps::Accumulator<S> > energy_acc( new pcps::EnergyAccumulator<S>(userinp, ham, wfn.overall_coefficient()) );
  ( dynamic_cast<pcps::EnergyAccumulator<S> *>(&*energy_acc) )->print() = true;
  walker.accumulators().push_back(energy_acc);

  const int nap = userinp.nparticles()/2;
  const int nas = userinp.nsites()/2;

  struct ConfFiller {
    static int fill(const int b, const int nas, int conf, std::vector<int> & conf_vec) {
      int count = 0;
      for (int i = b; i < 2*nas; i+=2) {
        conf_vec.at(i) = conf % 2;
        count += conf % 2;
        conf /= 2;
      }
      return count;
    }
  };

  // print a heading
  std::cout << boost::format("%14s  %14s") % "alpha conf" % "beta conf";
  ( dynamic_cast<pcps::EnergyAccumulator<S> *>(&*energy_acc) )->wfn_sampling() = true;
  ( dynamic_cast<pcps::EnergyAccumulator<S> *>(&*energy_acc) )->print_header();
  ( dynamic_cast<pcps::EnergyAccumulator<S> *>(&*energy_acc) )->wfn_sampling() = false;
  ( dynamic_cast<pcps::EnergyAccumulator<S> *>(&*energy_acc) )->print_header();
  std::cout << std::endl;

  //std::cout << boost::format("%14s  %14s  %18s  %18s  %24s  %24s  %24s  %24s  %24s  %24s")
  //             % "alpha conf"
  //             % "beta conf"
  //             % "<n|Psi>"
  //             % "<n|PsiG>"
  //             % "local 2rdm"
  //             % "<n| H_1 |Psi>"
  //             % "<Psi| L |n><n| R |Psi>"
  //             % "local te diagonal"
  //             % "local te half-diagonal"
  //             % "local te off-diagonal"
  //          << std::endl;

  // loop over alpha configurations
  const int max_a_conf = pcps::pow<int>(2, nas);
  for (int a_conf = 0; a_conf < max_a_conf; a_conf++) {

    // get the alpha configuration
    if ( ConfFiller::fill(0, nas, a_conf, walker.config().fock_config()) != nap )
      continue;

    // loop over beta configurations
    for (int b_conf = 0; b_conf < max_a_conf; b_conf++) {

      // get the beta configuration
      if ( ConfFiller::fill(1, nas, b_conf, walker.config().fock_config()) != nap )
        continue;

      walker.config().init_occ_unocc_site();

      // print the configuration info
      std::stringstream conf_str_a;
      for (int i = 0; i < 2*nas; i+=2)
        conf_str_a << walker.config().fock_config().at(i);
      std::cout << boost::format("%14s") % conf_str_a.str();
      std::stringstream conf_str_b;
      for (int i = 1; i < 2*nas; i+=2)
        conf_str_b << walker.config().fock_config().at(i);
      std::cout << boost::format("  %14s") % conf_str_b.str();

      //// initialize coefficient object for the given configuration
      //bool success = false;
      //for (int i = 0; i < 50 && !success; i++) {
      //  try {
      //    walker.config().generate_noise(userinp);
      //    walker.gf()->initialize(walker.config());
      //    success = ( std::abs( walker.gf()->get_value().get() ) > 1.0e-8 );
      //  } catch ( pcps::SingularMatrixException & e) {
      //    success = false;
      //  }
      //}

      // wavefunction guiding
      {
        ( dynamic_cast<pcps::EnergyAccumulator<S> *>(&*energy_acc) )->wfn_sampling() = true;
        ( dynamic_cast<pcps::RestPairingFunc<S> *>(&*walker.gf())  )->wfn_sampling() = true;
        bool success = false;
        try {
          walker.gf()->initialize(walker.config());
          success = ( std::abs( walker.gf()->get_value().get() ) > 1.0e-8 );
        } catch ( pcps::SingularMatrixException & e) {
          success = false;
        }
        if (success)
          walker.accumulators().at(0)->accumulate(userinp, 1.0, walker.config());
        else
          std::cout << boost::format("  %24s  %24s  %24s  %24s") % "" % "" % "" % "";
      }

      // rdm guiding
      {

        ( dynamic_cast<pcps::EnergyAccumulator<S> *>(&*energy_acc) )->wfn_sampling() = false;
        ( dynamic_cast<pcps::RestPairingFunc<S> *>(&*walker.gf())  )->wfn_sampling() = false;

        pcps::Configuration<S> resting(walker.config());

        // initialize coefficient object for the given configuration
        bool success = false;
        for (int b = 0; b < 2 && !success; b++) {

          for (int p = 0; p < nas && !success; p++) {

            if (walker.config().fock_config().at(2*p+b) == 0) continue;

            for (int q = 0; q < nas && !success; q++) {

              if ( p != q && walker.config().fock_config().at(2*q+b) == 1 ) continue;

              resting = walker.config();
              resting.fock_config().at(2*p+b) = 0;
              resting.fock_config().at(2*q+b) = 1;
              resting.init_occ_unocc_site();

              try {
                walker.gf()->initialize(resting); // start on resting
                if (p != q) {
                  pcps::ConfigDiff<S> diff(walker.config(), resting);
                  walker.gf()->get_value(walker.config(), resting, diff); // prepare to move to config
                  walker.gf()->update(walker.config(), resting, diff); // move to config
                }
                //success = ( std::abs( walker.gf()->get_value().get() ) > 1.0e-8 );
                success = std::abs( ( dynamic_cast<pcps::RestPairingFunc<S> *>(&*walker.gf())  )->init_value().get() ) > 1.0e-8;
              } catch ( pcps::SingularMatrixException & e) {
                success = false;
              }

            }

          }

        }

        if (success)
          walker.accumulators().at(0)->accumulate(userinp, 1.0, walker.config());
        else
          std::cout << boost::format("  %24s  %24s  %24s  %24s  %24s  %24s  %24s") % "" % "" % "" % "" % "" % "" % "";

      }

      // end the line
      std::cout << std::endl;

    }

  }

}

template <class S> void test_agp(const pcps::Input & userinp,
                                 pcps::Hamiltonian<S> & ham,
                                 pcps::Wavefunction<S> & wfn) {

  // create a function object for sampling energies
  struct GetEnergy {

    static double get(const pcps::Input & userinp,
                      pcps::Hamiltonian<S> & ham,
                      pcps::Wavefunction<S> & wfn,
                      const bool wfn_sampling,
                      double & pqqp_multiplier)
    {

      // create the walker
      pcps::VMC_Walker<S> walker(userinp);

      // set the guiding function
      walker.gf() = wfn.overall_coefficient();

      // create the energy accumulator
      boost::shared_ptr<pcps::Accumulator<S> > energy_acc( new pcps::EnergyAccumulator<S>(userinp, ham, wfn.overall_coefficient()) );

      // set the sampling type
      ( dynamic_cast<pcps::RestPairingFunc<S> *>(&*walker.gf())  )->wfn_sampling() = wfn_sampling;
      ( dynamic_cast<pcps::EnergyAccumulator<S> *>(&*energy_acc) )->wfn_sampling() = wfn_sampling;

      // get the [pq,pq] integrals for possible use in the guiding function
      ( dynamic_cast<pcps::RestPairingFunc<S> *>(&*walker.gf())  )->get_pqpq_ints(ham);

      // get the pqqp multiplier
      if (userinp.pqpq_guiding())
        ( dynamic_cast<pcps::EnergyAccumulator<S> *>(&*energy_acc) )->pqqp_multiplier() = pqqp_multiplier;

      // burn in
      pcps::sample(userinp, userinp.burn_in_length(), walker);

      // sample
      walker.accumulators().push_back(energy_acc);
      pcps::sample(userinp, userinp.sample_length(), walker, false);

      // save the pqqp multiplier
      if (wfn_sampling)
        pqqp_multiplier = ( dynamic_cast<pcps::EnergyAccumulator<S> *>(&*energy_acc) )->pqqp_multiplier();

      // print some statistics on the samplint
      if (pcps::mpi::rank() == 0)
        ( dynamic_cast<pcps::EnergyAccumulator<S> *>(&*energy_acc) )->print_statistics();

      // return the sampled energy
      return ( dynamic_cast<pcps::EnergyAccumulator<S>*>(energy_acc.get()) )->energy();

    }

  };

  // load the hamiltonian
  ham.from_archive("pcps_ham_dump.bin");

//  // load hamiltonian
//  if (pcps::mpi::rank() == 0) {
//    pcps::BinaryFileArchive ar("pcps_ham_dump.bin", "r");
//    ar >> ham;
//    for (size_t i = 0; i < ham.cholesky_vecs().size(); i++)
//      std::cout << boost::format(" %20.12f\n") % ham.cholesky_vecs()[i];
//    std::cout << std::endl;
//  } else {
//    pcps::BinaryArchive ar;
//    ar >> ham;
//  }

  //exact_sampling(userinp, ham, wfn);
  //return;

  //if ( pcps::mpi::rank() == 0 )
  //  std::cout << std::endl;

  // get a variable to hold the pqqp multiplier
  double pqqp_multiplier = 0.0;

  pcps::mpi::barrier();
  pcps::start_timer("wfn sampling");

  // get the one electron and semi-diagonal two electron energy contribution
  const double e_half_diag = GetEnergy::get(userinp, ham, wfn, true, pqqp_multiplier);

  pcps::mpi::barrier();
  pcps::stop_timer("wfn sampling");
  pcps::start_timer("rdm sampling");

  // get the off-diagonal two electron energy contribution
  const double e_off_diag = GetEnergy::get(userinp, ham, wfn, false, pqqp_multiplier);

  pcps::mpi::barrier();
  pcps::stop_timer("rdm sampling");

  // get the total energy
  const double e_total = e_half_diag + e_off_diag + ham.e_cor_nuc();

  //// save hamiltonian
  //if (pcps::mpi::rank() == 0) {
  //  pcps::BinaryFileArchive ar("pcps_ham_dump.bin");
  //  ar << ham;
  //}

  // print the energies
  if ( pcps::mpi::rank() == 0 ) {
    std::cout << boost::format("%24s  %24s  %24s  %24s  %24s")
                 % "pqqp multiplier"
                 % "1e and half-diagonal 2e"
                 % "off-diagonal 2e"
                 % "core/nuclear repulsion"
                 % "total"
              << std::endl;
    std::cout << boost::format("%24.12f  %24.12f  %24.12f  %24.12f  %24.12f")
                 % pqqp_multiplier
                 % e_half_diag
                 % e_off_diag
                 % ham.e_cor_nuc()
                 % e_total
              << std::endl;
    std::cout << std::endl;
  }

}

template void test_agp(const pcps::Input & userinp, pcps::Hamiltonian<double> &, pcps::Wavefunction<double> &);
template void test_agp(const pcps::Input & userinp, pcps::Hamiltonian<std::complex<double> > &, pcps::Wavefunction<std::complex<double> > &);
